# -*- coding: utf-8 -*-
"""
===============================
@Author     : Zuo WenTao
@Time       : 2024/7/15 9:10
@Description: pip install kafka-python
@Change     : 
@File       : KafKaCient.py
===============================
"""
import json

from loguru import logger
from kafka.admin import NewTopic
from kafka.errors import KafkaError
from kafka import KafkaProducer, KafkaConsumer, KafkaAdminClient


class KafkaClient:
    def __init__(self, brokers):
        self.brokers = [brokers]
        self.producer = None
        self.consumer = None
        self.admin_client = None

    def connect_producer(self):
        """ 连接 Kafka 生产者 """
        try:
            self.producer = KafkaProducer(
                bootstrap_servers=self.brokers,
                value_serializer=lambda v: json.dumps(v).encode('utf-8'),
                sasl_plain_username="admin",
                sasl_plain_password="#S9ZxYG3c#Wb"
            )
            logger.success(f"[Kafka] -- 生产者连接成功: {self.brokers}")
        except KafkaError as e:
            logger.error(f"[Kafka] -- 连接 Kafka 生产者失败: {self.brokers} - {e}")

    def connect_consumer(self, topic, group_id, auto_offset_reset='earliest'):
        """ 连接 Kafka 消费者 """
        try:
            self.consumer = KafkaConsumer(
                topic,
                bootstrap_servers=self.brokers,
                group_id=group_id,
                auto_offset_reset=auto_offset_reset,
                value_deserializer=lambda m: json.loads(m.decode('utf-8'))
            )
            logger.success(f"[Kafka] -- 消费者连接成功: {self.brokers}")
        except KafkaError as e:
            logger.error(f"[Kafka] -- 连接 Kafka 消费者失败: {self.brokers} - {e}")

    def connect_admin(self):
        """ 连接 Kafka 管理员客户端，用于管理主题 """
        try:
            self.admin_client = KafkaAdminClient(
                bootstrap_servers=self.brokers,
                sasl_plain_username="admin",
                sasl_plain_password="#S9ZxYG3c#Wb"
            )
            logger.success(f"[Kafka] -- 管理员客户端连接成功: {self.brokers}")
        except KafkaError as e:
            logger.error(f"[Kafka] -- 管理员客户端失败: {self.brokers} - {e}")

    def create_topic(self, topic_name, num_partitions=1, replication_factor=1):
        """ 创建新的 Kafka 主题 """
        if not self.admin_client:
            self.connect_admin()

        try:
            topic_list = [NewTopic(name=topic_name, num_partitions=num_partitions, replication_factor=replication_factor)]
            self.admin_client.create_topics(new_topics=topic_list, validate_only=False)
            logger.success(f"[Kafka] -- 主题创建成功: {topic_name}")
        except KafkaError as e:
            logger.error(f"[Kafka] -- 创建主题失败: {e}")

    def delete_topic(self, topic_name):
        """ 删除 KafKa 主题 """
        if not self.admin_client:
            self.connect_admin()

        try:
            self.admin_client.delete_topics(topics=[topic_name])
            logger.success(f"[Kafka] -- 主题删除成功: {topic_name}")
        except KafkaError as e:
            logger.error(f"[Kafka] -- 删除主题失败: {e}")

    def send_message(self, topic, message):
        """ 发送消息到 Kafka 主题 """
        if not self.producer:
            logger.error(f"[Kafka] -- 生产者未连接: {self.producer}")
            return

        try:
            self.producer.send(topic, message)
            self.producer.flush()
            logger.success(f"[Kafka] -- Topic: {topic} 消息发送成功: {message}")
        except KafkaError as e:
            logger.error(f"[Kafka] -- Topic: {topic} 发送消息失败: {e}")

    def update_message(self, topic, key: str, new_message: dict):
        """
        Kafka 不支持直接消息更新，这个函数本质事使用相同的键发送一条新的消息
        Args:
            topic: 主題名
            key: 需要修改的键
            new_message: 需要更新的信息

        Returns:

        """
        if not self.producer:
            logger.error(f"[Kafka] -- 生产者未连接: {self.producer}")
            return

        try:
            self.producer.send(topic, key=key, value=new_message)
            self.producer.flush()
            logger.success(f"[Kafka] -- Topic: {topic} 消息更新成功: {new_message}")
        except KafkaError as e:
            logger.error(f"[Kafka] -- Topic: {topic} 更新消息失败: {e}")

    def consume_messages(self, process_message_callback):
        if not self.consumer:
            logger.error(f"[Kafka] -- 消费者未连接: {self.producer}")
            return

        try:
            for message in self.consumer:
                logger.debug(f"[Kafka] -- 收到消息: {message.value}")
                process_message_callback(message.value)
        except KafkaError as e:
            logger.error(f"[Kafka] -- 消费消息失败: {e}")

    def close(self):
        if self.producer:
            self.producer.close()
        if self.consumer:
            self.consumer.close()
        if self.admin_client:
            self.admin_client.close()


# 使用示例:
if __name__ == "__main__":
    brokers = '192.168.44.182:9092'
    topic = 'my_topic'

    kafka_client = KafkaClient(brokers)
    kafka_client.connect_producer()
    # kafka_client.create_topic(topic_name='new_topic')
    message = {"key": "value"}
    kafka_client.send_message(topic, message)
    #
    # kafka_client.update_message(topic, key='key', new_message={"key": "new_value"})
    #
    # kafka_client.connect_consumer(topic, group_id='my_group')
    #
    #
    # def process_message(message):
    #     print(f"处理消息: {message}")
    #
    #
    # kafka_client.consume_messages(process_message_callback=process_message)
    #
    # kafka_client.delete_topic(topic_name='new_topic')
    #
    # kafka_client.close()
